"""批量且并行的生成数据集"""

import queue
import signal
import logging, time, os, shutil, sys
import multiprocessing as mp
from concurrent.futures import Future, ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Optional
from feasytools import time2str, ArgChecker

from main import simulate
from cspk_gen import generate_cspk
from trip_gen import generate_trips
from plot import plot_main


class SimConfig:
    def __init__(self, dname:str, respath:str, veh_cnt: int, buses: str, no_plg: str, log: str, seed:int,
            step_len:int, start_time:int, end_time:int, copy_cspk:bool, copy_trips:bool, plot_commands:str,
            no_parallel:bool=False):
        self.dname = dname
        self.respath = respath
        self.veh_cnt = veh_cnt
        self.buses = buses
        self.no_plg = no_plg
        self.log = log
        self.seed = seed
        self.step_len = step_len
        self.start_time = start_time
        self.end_time = end_time
        self.copy_cspk = copy_cspk
        self.copy_trips = copy_trips
        self.plot_commands = ArgChecker(plot_commands)
        self.no_parallel = no_parallel

def work(
    i: int,
    cfg: SimConfig,
    mpQ: Optional[queue.Queue[tuple[int, str]]] = None
):
    # 创建配置副本
    pname = f"_tmp_{cfg.dname}_{i:03}"
    shutil.copytree(cfg.dname, pname, dirs_exist_ok=True)
    shutil.move(pname + "/" + cfg.dname + ".sumocfg", pname + "/" + pname + ".sumocfg")
    shutil.move(pname + "/" + cfg.dname + ".net.xml", pname + "/" + pname + ".net.xml")
    shutil.move(pname + "/" + cfg.dname + ".taz.xml", pname + "/" + pname + ".taz.xml")
    shutil.move(pname + "/" + cfg.dname + ".plg.xml", pname + "/" + pname + ".plg.xml")
    shutil.move(pname + "/" + cfg.dname + ".grid.zip", pname + "/" + pname + ".grid.zip")
    # 计算充电桩数量
    slots = int(cfg.veh_cnt / (140 * 7))

    # 生成CS, PK
    if not cfg.copy_cspk: 
        generate_cspk(cfg.seed, pname, 0, cfg.buses.split(","), 1.0, 2 * slots, slots, silent = mpQ is not None)
    else: # 未指定母线则不重新生成CS, PK
        shutil.move(pname + "/" + cfg.dname + ".cs.xml", pname + "/" + pname + ".cs.xml")
        shutil.move(pname + "/" + cfg.dname + ".pk.xml", pname + "/" + pname + ".pk.xml")
    if mpQ: mpQ.put_nowait((i, "cspk"))
    
    # 生成Trips并模拟
    if not cfg.copy_trips:
        generate_trips(pname, cfg.seed, cfg.veh_cnt, 1.0, silent = mpQ is not None)
    else:
        shutil.move(pname + "/" + cfg.dname + ".veh.xml", pname + "/" + pname + ".veh.xml")
    if mpQ: mpQ.put_nowait((i, "trip"))

    outdir = f"results_{i}"
    shutil.move(pname + "/" + pname + ".grid.zip", pname + "/" + cfg.dname + ".grid.zip")
    shutil.move(pname + "/" + pname + ".net.xml", pname + "/" + cfg.dname + ".net.xml")
    shutil.move(pname + "/" + pname + ".taz.xml", pname + "/" + cfg.dname + ".taz.xml")
    simulate(
        cfg_dir=pname,
        traffic_step=cfg.step_len,
        start_time=cfg.start_time,
        end_time=cfg.end_time,
        outdir=outdir,
        no_plg=cfg.no_plg,
        result_fmt=".sdt",
        log=cfg.log,
        seed=cfg.seed,
        def_veh_cnt=cfg.veh_cnt,
        vb=None,
        copy=True,
        mpQ=mpQ,
        clntID=i,
    )
    
    # 生成图像
    plot_main(outdir, cfg.plot_commands, True)
    # 移动文件
    Path(cfg.respath).parent.mkdir(parents=True, exist_ok=True)
    shutil.move(outdir, cfg.respath)
    shutil.rmtree(pname)

def parallel_sim(
    parallel: int,
    configs: list[SimConfig]
):
    cpu_count = os.cpu_count()
    if cpu_count is None:
        print("Failed to get CPU count, set to 1")
        cpu_count = 1
    if parallel > cpu_count:
        print(
            "Your device are not able to handle so much parallel process efficiently, set to",
            cpu_count,
        )
        parallel = cpu_count
    pool = mp.Pool(parallel)
    mpQ = mp.Manager().Queue()
    skip_list = []
    exit_count = 0
    gen_trip_count = 0
    gen_cspk_count = 0
    
    def create_callback(i:int):
        def callback(e: BaseException):
            nonlocal exit_count
            exit_count += 1
            print(f"==========Error in Task {i}==========")
            logging.exception(e)
            print("=====================================")
        return callback

    st_time = time.time()
    non_paras:list[tuple[int,SimConfig]] = []
    for i, cfg in enumerate(configs):
        if Path(cfg.respath).exists():
            skip_list.append(i)
            continue
        if cfg.no_parallel:
            non_paras.append((i, cfg))
        else:
            pool.apply_async(work, (i, cfg, mpQ), error_callback=create_callback(i))
    pool.close()
        
    print("Skip list:", skip_list)
    rounds = len(configs)
    progs:list[float] = [0.] * rounds
    last_upd = 0
    real_rounds = rounds - len(skip_list) - len(non_paras)
    exit_count += len(skip_list) + len(non_paras)

    def eh(signum, frame):
        print("\nReceive SIGINT, exiting...")
        pool.terminate()
        pool.join()
        sys.exit()
    signal.signal(signal.SIGINT, eh)

    while exit_count < rounds:
        try:
            clnt, con = mpQ.get(timeout=1)
            clnt:int
            con:str
            if con == "exit":
                exit_count += 1
                progs[clnt] = 100
            elif con == "cspk":
                gen_cspk_count += 1
            elif con == "trip":
                gen_trip_count += 1
            else:
                progs[clnt] = float(con)
        except queue.Empty:
            pass
        if time.time() - last_upd > 1:
            if gen_cspk_count < real_rounds:
                cspk_info = f"CSPK: {gen_cspk_count}/{real_rounds}"
            else:
                cspk_info = ""
            if gen_trip_count < real_rounds:
                trip_info = f"Trips: {gen_trip_count}/{real_rounds}"
            else:
                trip_info = ""
            last_upd = time.time()
            elapsed = time2str(last_upd - st_time)
            avg_prog = sum(progs) / real_rounds
            eta = (
                time2str((last_upd - st_time) * (100 - avg_prog) / avg_prog)
                if avg_prog > 0
                else "N/A"
            )
            print(
                f"\rProgress: {avg_prog:.2f}% {cspk_info} {trip_info} Elapsed: {elapsed}  ETA: {eta}                ",
                end="",
            )

    print("\rFinishing plotting and file operations...                           ")
    pool.join()
    print(f"Parallel simulation done. Duration: {time2str(time.time()-st_time)}")
    if len(non_paras) > 0:
        print("Execute non-parallel tasks...")
        for i, cfg in non_paras:
            work(i, cfg)
        print("All done.")

HELP_STR = f"""
用法: python {sys.argv[0]} [-p <parallel_count>] [-r <results_root>] [-f <outier_configs>]
     [-d <config_directory> [-n <rounds>] [-c <veh_cnt>] [--copy-cspk] [--copy-trips]
     [-st <start_time>] [-et <end_time>] [-l <step_len>] [-b <buses>] [--plot <plot_command>]
     [-log <logging_level>] [-no-plg <disabled plugins>]] [-h/help]
    h/help: 显示本帮助信息
    p: 并行数, 默认为CPU核心数, 如果系统无法自动获取, 则必须指定
    r: 结果存储目录, 默认./dataset
    以下两组参数至少需要指定一组:
    组1:
        f: 外部参数文件
    组2:
        d: 配置文件目录
        n: 模拟次数, 默认50
        c: 每次模拟车辆数, 默认10000
        b: 使用的母线, 默认"B3,B8,B12,B16,B19,B24,B28,B32"
        log: 日志级别, 默认"cspk,gen,bus"
        no-plg: 不使用的插件, 默认为空
        st: 起始时间, 默认-1
        et: 结束时间, 默认-1
        l: 仿真步长, 默认10
        plot: 绘图指令, 默认"--load-accum"
        copy-cspk: 复制CS, PK文件, 默认False
        copy-trips: 复制Trips文件, 默认False
"""
def main():
    args = ArgChecker()
    if args.pop_bool("h") or args.pop_bool("help"):
        print(HELP_STR)
        exit()
    
    parallel = args.pop_int("p", os.cpu_count())
    root = args.pop_str("r", f"./dataset")
    from_file = args.pop_str("f","")
    if from_file == "":
        dname = args.pop_str("d")
        rounds = args.pop_int("n", 50)
        veh_cnt = args.pop_int("c", 10000)
        buses = args.pop_str("b", "B3,B8,B12,B16,B19,B24,B28,B32")
        noplg = args.pop_str("no-plg", "")
        log = args.pop_str("log", "cspk,gen,bus")
        step_len = args.pop_int("l", 10)
        start_time = args.pop_int("st", -1)
        end_time = args.pop_int("et", -1)
        copy_cspk = args.pop_bool("copy-cspk")
        copy_trips = args.pop_bool("copy-trips")
        plot_commands = args.pop_str("plot","--load-accum")
        configs = []
        for i in range(rounds):
            configs.append(SimConfig(dname, f"{root}/{i:03}", veh_cnt, buses, noplg,
                log, i, step_len, start_time, end_time, copy_cspk, copy_trips, plot_commands))
    else:
        with open(from_file, "r") as f:
            lines = f.readlines()
        configs = []
        for i, line in enumerate(lines):
            line = line.strip()
            if line == "" or line.startswith("#"):
                continue
            args = ArgChecker(line)
            configs.append(SimConfig(
                args.pop_str("d"), 
                f"{root}/{args.pop_str('o',f'{i:03}')}", 
                args.pop_int("c", 10000), 
                args.pop_str("b", ""), 
                args.pop_str("no-plg", ""),
                args.pop_str("log", "cspk,gen,bus"),
                args.pop_int("seed", i),
                args.pop_int("l", 10),
                args.pop_int("st", -1),
                args.pop_int("et", -1),
                args.pop_bool("copy-cspk"),
                args.pop_bool("copy-trips"),
                args.pop_str("plot","--load-accum"),
                args.pop_bool("no-parallel")
            ))
    parallel_sim(parallel, configs)

if __name__ == "__main__":
    main()
